import pytesseract
from PIL import Image
import cv2
import numpy as np


def find_basic_data_position_from_bytes(screenshot_data: bytes):
    """
    从字节数据中查找"基础数据"文字的位置信息
    :param screenshot_data: 屏幕截图的字节数据
    :return: "基础数据"文字的位置信息(x, y, width, height)或None
    """
    try:
        # 将字节数据转换为OpenCV图像
        nparr = np.frombuffer(screenshot_data, np.uint8)
        img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
        
        if img is None:
            raise ValueError("无法解码图像数据")

        # 转为 PIL 图像格式
        pil_img = Image.fromarray(img)
        # Linux下通常不需要指定tesseract路径，但如果需要可以取消下面一行的注释并设置正确的路径
        # pytesseract.pytesseract.tesseract_cmd = '/usr/bin/tesseract'

        # 使用 OCR 识别文字（支持中文）并获取位置信息
        # 使用页面分割模式6来更好地识别连续文本
        data = pytesseract.image_to_data(pil_img, lang='chi_sim+eng', output_type=pytesseract.Output.DICT, config='--psm 6')
        
        # 合并同一行上的文本以更好地识别连续文本
        lines = merge_text_lines(data)
        
        # 在合并后的行中查找目标文本"基础数据"
        for i, line in enumerate(lines):
            # 移除文本中的所有空格后再查找
            cleaned_text = ''.join(line['text'].split())
            # 只要包含"基础数据"就匹配
            if "基础数据" in cleaned_text:
                x, y, w, h = line['x'], line['y'], line['w'], line['h']
                print(f"找到'基础数据': (x={x}, y={y}, width={w}, height={h})")
                return x, y, w, h

        print("❌ 未找到'基础数据'")
        return None
    except Exception as e:
        print(f"OCR 识别失败: {e}")
        return None


def find_basic_data_position(image_path: str):
    """
    查找图像中"基础数据"文字的位置信息
    :param image_path: 图像路径
    :return: "基础数据"文字的位置信息(x, y, width, height)或None
    """
    # 读取图像
    img = cv2.imread(image_path)
    if img is None:
        raise ValueError(f"无法加载图像: {image_path}")

    # 转为 PIL 图像格式
    pil_img = Image.fromarray(img)
    # Linux下通常不需要指定tesseract路径，但如果需要可以取消下面一行的注释并设置正确的路径
    # pytesseract.pytesseract.tesseract_cmd = '/usr/bin/tesseract'

    # 使用 OCR 识别文字（支持中文）并获取位置信息
    try:
        # 使用页面分割模式6来更好地识别连续文本
        data = pytesseract.image_to_data(pil_img, lang='chi_sim+eng', output_type=pytesseract.Output.DICT, config='--psm 6')
        
        # 合并同一行上的文本以更好地识别连续文本
        lines = merge_text_lines(data)
        
        # 在合并后的行中查找目标文本"基础数据"
        for i, line in enumerate(lines):
            # 移除文本中的所有空格后再查找
            cleaned_text = ''.join(line['text'].split())
            print(f"正在处理行: '{cleaned_text}'")
            # # 只要包含"基础数据"就匹配
            # if "基础数据" in cleaned_text:
            #     x, y, w, h = line['x'], line['y'], line['w'], line['h']
            #     print(f"找到'基础数据': (x={x}, y={y}, width={w}, height={h})")
            #     return x, y, w, h

        # print("❌ 未找到'基础数据'")
        return None
    except Exception as e:
        print(f"OCR 识别失败: {e}")
        return None


def merge_text_lines(data, y_tolerance=10):
    """
    合并同一行上的文本框
    
    :param data: pytesseract 返回的数据字典
    :param y_tolerance: Y坐标容差，用于判断是否在同一行
    :return: 合并后的行列表
    """
    lines = []
    n_boxes = len(data['level'])
    
    # 提取有效的文本框
    boxes = []
    for i in range(n_boxes):
        text = data['text'][i].strip()
        # 去除空字符串但仍显示空白字符
        if len(text) > 0:
            box = {
                'text': text,
                'x': data['left'][i],
                'y': data['top'][i],
                'w': data['width'][i],
                'h': data['height'][i]
            }
            boxes.append(box)
    
    # 根据Y坐标对文本框进行分组
    if not boxes:
        return lines
        
    # 按Y坐标排序
    boxes.sort(key=lambda b: b['y'])
    
    # 分组同一行的文本框
    current_line = [boxes[0]]
    for i in range(1, len(boxes)):
        # 如果当前文本框与前一个文本框在同一行，则添加到当前行
        if abs(boxes[i]['y'] - boxes[i-1]['y']) <= y_tolerance:
            current_line.append(boxes[i])
        else:
            # 否则，结束当前行并将文本框按X坐标排序后合并
            current_line.sort(key=lambda b: b['x'])
            merged_text = ' '.join([b['text'] for b in current_line])
            
            # 计算合并后的边界框
            min_x = min([b['x'] for b in current_line])
            max_x = max([b['x'] + b['w'] for b in current_line])
            min_y = min([b['y'] for b in current_line])
            max_y = max([b['y'] + b['h'] for b in current_line])
            
            lines.append({
                'text': merged_text,
                'x': min_x,
                'y': min_y,
                'w': max_x - min_x,
                'h': max_y - min_y
            })
            
            # 开始新行
            current_line = [boxes[i]]
    
    # 处理最后一行
    if current_line:
        current_line.sort(key=lambda b: b['x'])
        merged_text = ' '.join([b['text'] for b in current_line])
        
        # 计算合并后的边界框
        min_x = min([b['x'] for b in current_line])
        max_x = max([b['x'] + b['w'] for b in current_line])
        min_y = min([b['y'] for b in current_line])
        max_y = max([b['y'] + b['h'] for b in current_line])
        
        lines.append({
            'text': merged_text,
            'x': min_x,
            'y': min_y,
            'w': max_x - min_x,
            'h': max_y - min_y
        })
    
    return lines


# 示例调用
if __name__ == '__main__':
    # Linux下的示例路径
    screenshot_path = r'D:\PycharmProjects\auto-script\screentshot\1.png'
    position = find_basic_data_position(screenshot_path)
    if position:
        print(f"位置信息: {position}")
    else:
        print("未找到'基础数据'")